import socket
import threading
import requests
import random
from qkdAppPacket import QKDAppPacket
from qkdBuffer import QKDBuffer
from qkdKeyGen import QKDKeyGen
from qkdKeyManageSystem import QKDKeyManagerSystem

class QKDApp:
    def __init__(self, role: str, local_address: tuple = None, remote_address: tuple = None):
        self.role = role  # 角色 'master' 或 'slave'
        self.local_address = local_address
        self.remote_address = remote_address
        self.sender_socket = None
        self.msg_id = random.randint(10000, 50000)  # 当前发送的数据包 ID，随机生成一个五位数

        if self.role == 'master' and (not self.local_address or not self.remote_address):
            raise ValueError("Master角色必须指定本地地址 (local_address) 和远程地址 (remote_address)。")
        if self.role == 'slave' and not self.local_address:
            raise ValueError("Slave角色必须指定本地地址 (local_address)，但不应指定远程地址 (remote_address)。")

    def _send(self, packet: QKDAppPacket):
        if self.sender_socket is None:
            self.sender_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.sender_socket.bind(self.local_address)

        # debug
        packet.print_packet()
        # 创建数据包
        packed_data = packet.pack()
        print(f"sender send packed_data:{packed_data}")
        # 发送数据包
        self.sender_socket.sendto(packed_data, self.remote_address)
        print("数据包已发送")

    def send_with_message(self, msg: str):
        if self.role != 'master':
            raise Exception("只有主节点可以发送数据包。")

        msg_bytes = msg.encode('utf-8')
        key_id, key_value = self.request_key()
        key_length = len(key_value)

        offset = 0
        while offset < len(msg_bytes):
            # 截取与密钥相同长度的消息片段，如果消息长度小于密钥长度，截取相同长度的密钥
            current_payload = msg_bytes[offset:offset + key_length]
            current_key = key_value[:len(current_payload)]
            encrypted_payload = self.xor_encrypt(current_payload, current_key)
            packet = QKDAppPacket(
                msg_id=self.msg_id,
                payload=encrypted_payload,
                encryption_key_id=key_id
            )
            # 计算数据包总长度
            packet.length = QKDAppPacket.header_length + len(packet.payload)
            # 使用内部函数发送数据包
            self._send(packet)

            # 递增消息 ID
            self.msg_id += 1
            offset += key_length

    def sendTest(self):
        print(f"start send test")
        if self.role != 'master':
            raise Exception("只有主节点可以发送数据包。")

        while True:
            time.sleep(1)
            message = f"Test message with ID {self.msg_id}"
            print(f"[DEBUG] Sending test message: {message}")
            self.send_with_message(message)

    def receive(self):
        if self.role != 'slave':
            raise Exception("接收失败：只有从节点可以接收数据包。")

        receiver_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        receiver_socket.bind(self.local_address)
        print("接收方正在监听...")

        try:
            while True:
                # 接收数据包
                data, addr = receiver_socket.recvfrom(4096)
                print("[DEBUG] 接收到来自", addr, "的数据包")

                # 解包数据包
                packet = QKDAppPacket.unpack(data)
                print(f"[DEBUG] 接收到的数据包:\n 长度: {packet.length}\n 消息 ID: {packet.msg_id}\n 加密算法: {packet.encryption_algo}\n 认证算法: {packet.auth_algo}\n 加密密钥 ID: {packet.encryption_key_id}\n 认证密钥 ID: {packet.auth_key_id}\n 认证标签: {packet.auth_tag}\n 负载: {packet.payload}")

                # 从 KMS 请求用于解密的密钥
                key_value = self.request_key_by_id(packet.encryption_key_id)
                current_key = key_value[:len(packet.payload)]
                decrypted_payload = self.xor_decrypt(packet.payload, current_key)
                print(f"[DEBUG] 解密后的负载: {decrypted_payload.decode('utf-8')}")
        
        except KeyboardInterrupt:
            print("接收线程已被手动终止")

        finally:
            receiver_socket.close()


    def request_key(self) -> tuple:
        # 从 KMS 请求密钥
        url = f"http://{self.local_address[0]}:8080/api/keys/request_keys?num_keys=1&ip_address={self.local_address[0]}"
        response = requests.get(url)
        if response.status_code == 200:
            key_info = response.json()[0]
            print(f"debug in request_key:key_info{key_info}")
            return key_info['key_id'], bytes.fromhex(key_info['key_value'])
        else:
            raise Exception(f"从 KMS 请求密钥失败，HTTP 状态码: {response.status_code}，请检查 KMS 服务状态。")

    def request_key_by_id(self, key_id: int) -> bytes:
        # 通过密钥 ID 从 KMS 请求密钥
        print(f"debug in request_key_by_id in local_address:{self.local_address[0]}")
        # for test
        # self.query_keys_for_ip(self.local_address[0]) 
        url = f"http://{self.local_address[0]}:8080/api/keys/get_key_with_ids?ip_address={self.local_address[0]}&key_IDs={key_id}"
        response = requests.get(url)
        print(f"request_key_by_id response:responseJson{response.json()}")
        if response.status_code == 200:
            key_info = response.json()[0]
            print(f"request_key_by_id response:key_info{key_info}")
            return bytes.fromhex(key_info['key_value'])
        else:
            raise Exception(f"通过密钥 ID 从 KMS 请求密钥失败，HTTP 状态码: {response.status_code}，请检查 KMS 服务状态。")
        
    def query_keys_for_ip(self, ip_address: str):
        # 查询某个 QKDBuffer 当前的秘钥
        url = f"http://{ip_address}:8080/api/keys/query_keys_for_ip?ip_address={ip_address}"
        response = requests.get(url)
        if response.status_code == 200:
            key_ids = response.json().get("key_ids", [])
            print(f"[DEBUG] 当前 QKDBuffer ({ip_address}) 中的密钥 ID 列表: {key_ids}")
        else:
            raise Exception(f"查询 QKDBuffer 中的密钥失败，HTTP 状态码: {response.status_code}，请检查 KMS 服务状态。")


    @staticmethod
    def xor_encrypt(data: bytes, key: bytes) -> bytes:
        return bytes(a ^ b for a, b in zip(data, key))

    @staticmethod
    def xor_decrypt(data: bytes, key: bytes) -> bytes:
        return bytes(a ^ b for a, b in zip(data, key))

if __name__ == "__main__":
    import time
    # 创建KMS并添加QKDBuffer
    kms1 = QKDKeyManagerSystem(kms_id="KMS001", max_keys_per_request=5, min_key_size=128, max_key_size=256, kms_hostname="localhost")
    kms2 = QKDKeyManagerSystem(kms_id="KMS002", max_keys_per_request=5, min_key_size=128, max_key_size=256, kms_hostname="localhost")
    buffer1 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="192.168.1.1")
    buffer2 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="192.168.1.2")
    time.sleep(1)
    kms1.add_buffer(buffer1)
    kms2.add_buffer(buffer2)

    # 启动KMS服务器
    threading.Thread(target=kms1.run_server, daemon=True).start()
    threading.Thread(target=kms2.run_server, daemon=True).start()
    time.sleep(1)  # 增加等待时间以确保KMS服务器已完全启动
    # 通过调用 /status API 检查 KMS 服务器是否启动成功
    try:
        response = requests.get("http://192.168.1.1:8080/api/keys/status")
        if response.status_code == 200:
            print("KMS1服务器启动成功")
        else:
            raise Exception("KMS服务器启动失败，状态码：" + str(response.status_code))
    except requests.exceptions.RequestException as e:
        print(f"无法连接到 KMS 服务器: {e}")
        exit(1)
    try:
        response = requests.get("http://192.168.1.2:8080/api/keys/status")
        if response.status_code == 200:
            print("KMS2服务器启动成功")
        else:
            raise Exception("KMS服务器启动失败，状态码：" + str(response.status_code))
    except requests.exceptions.RequestException as e:
        print(f"无法连接到 KMS 服务器: {e}")
        exit(1)
    # 启动QKDKeyGen在两个QKDBuffer之间创建密钥
    key_gen = QKDKeyGen(key_size=256, key_rate=5, packet_size=1024, data_rate=100, max_sifting_packets=10, master_address="192.168.1.1", slave_address="192.168.1.2")
    threading.Thread(target=key_gen.start_communication, daemon=True).start()
    time.sleep(1)
    # 创建QKDApp进行通信
    qkd_app_master = QKDApp(role="master", local_address=('192.168.1.1', 54321), remote_address=('192.168.1.2', 12345))
    qkd_app_slave = QKDApp(role="slave", local_address=('192.168.1.2', 12345))

    # 测试时，分别在独立线程中运行发送方和接收方
    receiver_thread = threading.Thread(target=qkd_app_slave.receive)
    sender_thread = threading.Thread(target=qkd_app_master.sendTest)

    receiver_thread.start()
    sender_thread.start()

    sender_thread.join()
    receiver_thread.join()
